"""
任务二配置文件

配置深度学习模型的超参数和训练设置

作者：数学建模团队
版本：2.0
"""

# 数据配置
DATA_CONFIG = {
    'csv_path': '../task1/task1_fixed_features_20250921_153924.csv',
    'test_size': 0.2,
    'val_size': 0.2,  # 从训练集中划分验证集的比例
    'random_state': 42,
    'batch_size': 32,
    'shuffle': True
}

# 模型配置
MODEL_CONFIGS = {
    'CNN': {
        'model_type': 'cnn',
        'dropout_rate': 0.5,
        'description': '一维卷积神经网络'
    },
    'ResNet': {
        'model_type': 'resnet',
        'dropout_rate': 0.5,
        'description': '一维残差网络'
    },
    'LSTM': {
        'model_type': 'lstm',
        'hidden_size': 128,
        'num_layers': 2,
        'dropout_rate': 0.5,
        'description': '长短期记忆网络'
    },
    'Hybrid': {
        'model_type': 'hybrid',
        'dropout_rate': 0.5,
        'description': 'CNN+LSTM混合模型'
    },
    'Attention': {
        'model_type': 'attention',
        'dropout_rate': 0.5,
        'description': '注意力机制模型'
    }
}

# 训练配置
TRAINING_CONFIG = {
    'epochs': 100,
    'learning_rate': 0.001,
    'patience': 15,  # 早停耐心值
    'weight_decay': 1e-5,
    'optimizer': 'Adam',
    'scheduler': 'ReduceLROnPlateau',
    'scheduler_params': {
        'mode': 'min',
        'factor': 0.5,
        'patience': 5,
        'verbose': True
    }
}

# 设备配置
DEVICE_CONFIG = {
    'auto_detect': True,  # 自动检测GPU
    'device': None,  # 如果auto_detect=False，手动指定设备
    'cuda_available': False  # 运行时检测
}

# 集成学习配置
ENSEMBLE_CONFIG = {
    'method': 'voting',  # 'voting' 或 'averaging'
    'top_k_models': 3,  # 选择性能最好的k个模型
    'weights': None  # 如果为None，使用等权重
}

# 可视化配置
VISUALIZATION_CONFIG = {
    'figure_size': (15, 10),
    'dpi': 300,
    'save_format': 'png',
    'font_size': 12,
    'color_palette': 'Set2'
}

# 文件路径配置
FILE_PATHS = {
    'models_dir': 'task2/',
    'results_dir': 'task2/',
    'plots_dir': 'task2/',
    'reports_dir': 'task2/'
}

# 类别配置
CLASS_CONFIG = {
    'class_names': ['Ball', 'Inner Race', 'Outer Race', 'Normal'],
    'class_mapping': {
        'Ball': 0,
        'Inner Race': 1, 
        'Outer Race': 2,
        'Normal': 3
    }
}

# 特征配置
FEATURE_CONFIG = {
    'feature_types': {
        '时域特征': ['mean', 'std', 'max', 'min', 'peak_to_peak', 'rms', 
                   'skewness', 'kurtosis', 'crest_factor', 'shape_factor', 
                   'impulse_factor', 'clearance_factor'],
        '频域特征': ['freq_mean', 'freq_std', 'freq_max', 'freq_centroid', 
                   'band_1_energy', 'band_2_energy', 'band_3_energy'],
        '时频域特征': ['wavelet_0_energy', 'wavelet_0_std', 'wavelet_1_energy', 
                     'wavelet_1_std', 'wavelet_2_energy', 'wavelet_2_std',
                     'wavelet_3_energy', 'wavelet_3_std', 'wavelet_4_energy', 
                     'wavelet_4_std', 'stft_mean', 'stft_std'],
        '故障频率特征': ['bpfo_harmonic_1_energy', 'bpfo_harmonic_2_energy', 
                       'bpfo_harmonic_3_energy', 'bpfi_harmonic_1_energy', 
                       'bpfi_harmonic_2_energy', 'bpfi_harmonic_3_energy',
                       'bsf_harmonic_1_energy', 'bsf_harmonic_2_energy', 
                       'bsf_harmonic_3_energy', 'ftf_harmonic_1_energy', 
                       'ftf_harmonic_2_energy', 'ftf_harmonic_3_energy'],
        '非线性特征': ['approx_entropy', 'sample_entropy', 'permutation_entropy']
    }
}

# 评估指标配置
EVALUATION_CONFIG = {
    'metrics': ['accuracy', 'precision', 'recall', 'f1_score'],
    'average_method': 'weighted',  # 'micro', 'macro', 'weighted'
    'confusion_matrix': True,
    'classification_report': True
}

# 日志配置
LOGGING_CONFIG = {
    'level': 'INFO',
    'format': '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    'log_file': 'task2/task2_training.log'
}

# 随机种子配置
RANDOM_SEED_CONFIG = {
    'numpy_seed': 42,
    'torch_seed': 42,
    'python_seed': 42
}

def get_config():
    """获取完整配置"""
    return {
        'data': DATA_CONFIG,
        'models': MODEL_CONFIGS,
        'training': TRAINING_CONFIG,
        'device': DEVICE_CONFIG,
        'ensemble': ENSEMBLE_CONFIG,
        'visualization': VISUALIZATION_CONFIG,
        'files': FILE_PATHS,
        'classes': CLASS_CONFIG,
        'features': FEATURE_CONFIG,
        'evaluation': EVALUATION_CONFIG,
        'logging': LOGGING_CONFIG,
        'random_seed': RANDOM_SEED_CONFIG
    }

def print_config():
    """打印配置信息"""
    config = get_config()
    print("任务二配置信息:")
    print("=" * 50)
    
    for section, params in config.items():
        print(f"\n{section.upper()}:")
        for key, value in params.items():
            print(f"  {key}: {value}")

if __name__ == "__main__":
    print_config()

